import numpy as np

def predict_action(state):
    """
    输入:
        state: np.ndarray, shape (b, l, 20)
    输出:
        action: np.ndarray, shape (b, 1), 元素为 [-1, 0, 1]
    """
    b, l, c = state.shape
    assert c == 20, "通道数必须是 20"

    # 指标索引映射
    IDX = {
        "ema_12": 8,
        "ema_26": 9,
        "macd": 10,
        "macd_signal": 11,
        "rsi": 12,
        "obv": 18
    }

    # 取当前和前一时刻
    macd = state[:, -1, IDX["macd"]]
    macd_prev = state[:, -2, IDX["macd"]]
    macd_signal = state[:, -1, IDX["macd_signal"]]
    macd_signal_prev = state[:, -2, IDX["macd_signal"]]

    ema_12 = state[:, -1, IDX["ema_12"]]
    ema_26 = state[:, -1, IDX["ema_26"]]

    rsi = state[:, -1, IDX["rsi"]]
    rsi_prev = state[:, -2, IDX["rsi"]]

    obv = state[:, -1, IDX["obv"]]
    obv_prev = state[:, -2, IDX["obv"]]

    # ==== 多头条件 ====
    macd_gold = (macd > macd_signal) & (macd_prev <= macd_signal_prev)
    ema_bull = ema_12 > ema_26
    rsi_up = (rsi < 70) & (rsi > rsi_prev)
    obv_up = obv > obv_prev
    long_signal = macd_gold & ema_bull & rsi_up & obv_up

    # ==== 空头条件 ====
    macd_dead = (macd < macd_signal) & (macd_prev >= macd_signal_prev)
    ema_bear = ema_12 < ema_26
    rsi_down = (rsi > 30) & (rsi < rsi_prev)
    obv_down = obv < obv_prev
    short_signal = macd_dead & ema_bear & rsi_down & obv_down

    # 构造动作
    action = np.zeros(b, dtype=np.int32)
    action[long_signal] = 1
    action[short_signal] = -1

    return action.reshape(-1, 1)
 